package goodash

import (
	"fmt"
	"math/rand/v2"
)

type IWeightAble interface {
	Weight() int
}

//复杂的随机算法

func RandomByWeightArray(weights []int) (int, error) {
	if len(weights) == 0 {
		return -1, fmt.Errorf("weights is empty")
	}
	totalWeight := Reduce(weights, func(agg int, item int, _ int) int {
		return agg + item
	}, 0)
	if totalWeight <= 0 {
		return -1, fmt.Errorf("totalWeight is less than or equal to 0")
	}
	choose := rand.IntN(totalWeight)
	for i, ele := range weights {
		choose -= ele
		if choose <= 0 {
			return i, nil
		}
	}
	return -1, fmt.Errorf("weights is error")
}

func RandomByWeightArray1(weights []int, totalWeight int) (int, error) {
	if len(weights) == 0 {
		return -1, fmt.Errorf("weights is empty")
	}
	//totalWeight := Reduce(weights, func(agg int, item int, _ int) int {
	//	return agg + item
	//}, 0)
	if totalWeight <= 0 {
		return -1, fmt.Errorf("totalWeight is less than or equal to 0")
	}
	choose := rand.IntN(totalWeight)
	for i, ele := range weights {
		choose -= ele
		if choose <= 0 {
			return i, nil
		}
	}
	return -1, fmt.Errorf("weights is error")
}
func RandomByWeightAbleArray[T IWeightAble](weights []T) (int, T, error) {
	if len(weights) == 0 {
		var zero T
		return -1, zero, fmt.Errorf("weights is empty")
	}
	totalWeight := Reduce(weights, func(agg int, item T, _ int) int {
		return agg + item.Weight()
	}, 0)
	if totalWeight <= 0 {
		var zero T
		return -1, zero, fmt.Errorf("totalWeight is less than or equal to 0")
	}
	choose := rand.IntN(totalWeight)
	for i, ele := range weights {
		choose -= ele.Weight()
		if choose <= 0 {
			return i, ele, nil
		}
	}
	var zero T
	return -1, zero, fmt.Errorf("weights is error")
}

// [[100,20],[300,50]
// weights = [20,50]
func RandomByWeightArrays(weights [][]int, weightIndex int) (int, []int, error) {
	if len(weights) == 0 {
		return -1, nil, fmt.Errorf("weights is empty")
	}
	totalWeight := Reduce(weights, func(agg int, item []int, _ int) int {
		return agg + item[weightIndex]
	}, 0)
	if totalWeight <= 0 {
		return -1, nil, fmt.Errorf("totalWeight is less than or equal to 0")
	}
	choose := rand.IntN(totalWeight)
	for i, ele := range weights {
		choose -= ele[weightIndex]
		if choose <= 0 {
			return i, ele, nil
		}
	}
	return -1, nil, fmt.Errorf("weights is error")
}

func RandomSizeByWeightArray(weights []int, size int, canRepeat bool) ([]int, error) {
	if len(weights) == 0 {
		return nil, fmt.Errorf("weights is empty")
	}
	if size <= 0 {
		return nil, fmt.Errorf("size is less than or equal to 0")
	}
	if size > len(weights) && !canRepeat {
		return nil, fmt.Errorf("size is greater than weights length")
	}

	var allIndexes []int
	totalWeight := Reduce(weights, func(agg int, item int, i int) int {
		if !canRepeat && item > 0 {
			allIndexes = append(allIndexes, i)
		}
		return agg + item
	}, 0)
	if totalWeight <= 0 {
		return nil, fmt.Errorf("totalWeight is less than or equal to 0")
	}
	if !canRepeat && size >= len(allIndexes) {
		return allIndexes, nil
	}

	var resultIndexes []int

	if !canRepeat {
		for i := 0; i < size; i++ {
			j, er := RandomByWeightArray1(weights, totalWeight)
			if er != nil {
				return nil, er
			}
			totalWeight -= weights[j]
			resultIndexes = append(resultIndexes, j)
			// 删除 j
			weights = append(weights[:j], weights[j+1:]...)
		}

	} else {
		for i := 0; i < size; i++ {
			choose := rand.IntN(totalWeight)
			for j, ele := range weights {
				choose -= ele
				if choose <= 0 {
					resultIndexes = append(resultIndexes, j)
				}
			}
		}
	}

	return resultIndexes, nil
}
